import logging
import math
import gym
from gym import spaces
from gym.utils import seeding
import numpy as np
import random
import os
import torch
import torch.nn as nn

from queue import Queue

# DIR = "data_test"
DIR = "data_train"

# backup原因
# 要调整state至18维，涉及改动较大

logger = logging.getLogger(__name__)

class DataSet():
    def __init__(self):
        # 二维向量，横轴是分片，纵轴是时隙
        self.bitratelevel = []
        self.bitrateInfo = []
        self.buffer = []
        self.FovRecord = []
        self.rebuffer_time = []
        self.throughput = []
        self.pause_total = []


    def Fetch_data_from_file(self,filepath,filename):
        pos = []
        file = filepath+filename
        with open(file, 'r') as file_to_read:
            while True:
                lines = file_to_read.readline()  # 整行读取数据
                if not lines:
                    break
                p_tmp = [float(i) for i in lines.split(',')]  # 将整行数据分割处理，如果分割符是空格，括号里就不用传入参数，如果是逗号， 则传入‘，'字符。
                pos.append(p_tmp)  # 添加新读取的数据
            pos = np.array(pos)
        if filename == "bitrateLevel.txt":
            self.bitratelevel = pos
        elif filename == "bitrateInfo.txt":
            self.bitrateInfo = pos
        elif filename == "throughput.txt":
            self.throughput = pos
        elif filename == "buffer.txt":
            self.buffer = pos
        elif filename == "FovRecord.txt":
            self.FovRecord = pos
        elif filename == "pause_total.txt":
            self.pause_total = pos
        elif filename == "rebuffer_time.txt":
            self.rebuffer_time = pos




class TiledVideo(gym.Env):
    metadata = {
        'render.modes': ['human', 'rgb_array'],
        'video.frames_per_second' : 50
    }

    def __init__(self):

        self.video = np.array([0,1,2,3,4,5]) # 2*2 tiled  不懂 video如何
        self.action_space = spaces.Discrete(4*4*4*4*4*4) # 4^6 每一个tile可以选择4种码率
        #self.action_space = spaces.MultiDiscrete(4 * np.ones((4,)))
        self.observation_space = spaces.Discrete(6) # 6
        # self.observation_space.n = 14
        #self.observation_space = spaces.MultiDiscrete(4 * np.ones((4,)))
        self.rng = np.random.RandomState(1234) # 1234是随机数种子
        self.state = [0,0,0,0,0,0] #
        self.ep_steps = 0 # 
        self.video_slot = 0
        self.record_bitratelevel = []
        self.dirs = [] # 存文件数据文件名的
        # self.tile_adjacent_relation = {0:[1,2],
        #                                1:[0,3],
        #                                2:[0,3],
        #                                3:[1,2]}
        self.tile_adjacent_relation = {0:[2,3,4,5],
                                       1:[2,3,4,5],
                                       2:[0,1,4,5],
                                       3:[0,1,4,5],
                                       4:[0,1,2,3],
                                       5:[0,1,2,3]}  # tile与其他tile的关联程度，是自行设计的
        self.buffer = [0,0,0,0,0,0]  # 每一个tile有一个buffer
        self.last_bitrate = [0,0,0,0,0,0] # [[0,0,0,0]] # 4个tile 1个时隙  上一时刻的比特率
        #self.last_throughout = [] 
        self.rebuffer_event = [] # rebuffer事件

        # 读取数据集信息
        self.dataset = DataSet() # 数据存储类

        # self.dataset.Fetch_data_from_file("data/2021-10-28-1/", "bitrateLevel.txt")
        # self.dataset.Fetch_data_from_file("data/2021-10-28-1/", "bitrateInfo.txt")
        # self.dataset.Fetch_data_from_file("data/2021-10-28-1/", "throughput.txt")
        # self.dataset.Fetch_data_from_file("data/2021-10-28-1/", "buffer.txt")
        # self.dataset.Fetch_data_from_file("data/2021-10-28-1/", "FovRecord.txt")
        # self.dataset.Fetch_data_from_file("data/2021-10-28-1/", "pause_total.txt")
        # self.dataset.Fetch_data_from_file("data/2021-10-28-1/", "rebuffer_time.txt")

        # self.chunk_length = len(self.dataset.buffer) # 视频长度为

        # record
        self.last_bitrate = [] # 上一个帧的比特率
        self.last_throughout = []

    def load_data(self):
        # self.dirs = os.listdir("D:\Project_Python\ABR\option-critic-pytorch-master\data_train")
        self.dirs = os.listdir("./"+DIR)
        
        dir_path = DIR + "/" + random.choice(self.dirs) + "/" # 一次随机选一个观看记录
        print(dir_path)
        self.dataset.Fetch_data_from_file(dir_path, "bitrateLevel.txt")
        self.dataset.Fetch_data_from_file(dir_path, "bitrateInfo.txt")
        self.dataset.Fetch_data_from_file(dir_path, "throughput.txt")
        self.dataset.Fetch_data_from_file(dir_path, "buffer.txt")
        self.dataset.Fetch_data_from_file(dir_path, "FovRecord.txt")
        self.dataset.Fetch_data_from_file(dir_path, "pause_total.txt")
        self.dataset.Fetch_data_from_file(dir_path, "rebuffer_time.txt")

    def seed(self, seed=None):
        return self._seed(seed)

    def _seed(self, seed=None):
        self.np_random, seed = seeding.np_random(seed)
        return [seed]

    def reset(self):
        self.state = [0, 0, 0, 0, 0, 0]
        self.buffer = [0,0,0,0,0,0]
        # self.last_bitrate = []
        self.last_bitrate = [] # [[0,0,0,0]] # 4个tile 1个时隙
        self.last_throughout = []
        self.rebuffer_event = []
        self.ep_steps = 0 # 没用上啊，main_ABR的程序里自己有记录
        self.record_bitratelevel = []  # 归零更新
        self.load_data()  # 在这里读文件
        # self.video_slot = 0  ## reset有两种情况，一个是视频播完，另一个是步长太大，所以不能在这个reset里归零video_slot
        return self.state

    def ten_to_4(self, x): # 10转4进制
        output = ""
        ret = x // 4
        leave = x % 4
        if ret == 0:
            return str(leave)
        xxx = self.ten_to_4(ret)
        output = xxx + str(leave)
        return output

    def Dimensionality_reduction(self,FoV_4_6):
        FoV_6 = [0,0,0,0,0,0]
        for i in range(6):
            FoV_6[i] = 0.5*FoV_4_6[0][i] + FoV_4_6[1][i] + FoV_4_6[2][i]

        return FoV_6

    def step(self, action, current_option): # 将已经判断出的action current_option送入，获取reward？

        # action = action.argmax()
        rebuffer_record = 0
        FoV = self.dataset.FovRecord
        throughput = self.dataset.throughput
        done = False

        # state_rate_mapping = {0:777451,
        #                       1:2769991,
        #                       2:6430841,
        #                       3:116736619}
        state_rate_mapping = {0: 100000000,
                              1: 300000000,
                              2: 500000000,
                              3: 700000000}

        FoV_Adjacent = {0:[2,3,4,5],
                        1:[2,3,4,5],
                        2:[0,1,4,5],
                        3:[0,1,4,5],
                        4:[0,1,2,3],
                        5:[0,1,2,3]}

        action = [action[i].detach().item() for i in range(6)]
        # for i in range(6):
        #     self.state[i] = action[i]

        self.record_bitratelevel.append(self.state) # 记录之前的码率选择
        # FoV_4_6 = []
        
        # for i in range(4): # 将现在 1s后 2s后 3s后放进去
        #     FoV_4_6.append(FoV[i])
        # FoV_state = self.Dimensionality_reduction(FoV_4_6)
        # 计算buffer，顺便记录rebuffer次数和值
        # for i in range(6):
        #     self.buffer[i] += throughput[self.video_slot+1][i]/state_rate_mapping[action[i]] - 1
        #     if self.buffer[i] < 0:  # 不够0就触发rebuffer事件，记录差的数值，是负数，在奖励里直接加就好
        #         rebuffer_record += self.buffer[i]
        #         self.buffer[i] = 0
        rebuffer_event = 0
        for i in range(6):
            if self.buffer[i] < 1:
                self.buffer[i] += throughput[self.video_slot][i] / state_rate_mapping[action[i]]
                if self.buffer[i] < 1:
                    rebuffer_event += 1 - self.buffer[i]
                    self.buffer[i] = 1
            else:
                self.buffer[i] = 1 + throughput[self.video_slot][i] / state_rate_mapping[action[i]]

            self.buffer[i] = self.buffer[i] - 1  # 模拟buffer过程
        # 现在计算出了每一个slot的 buffer，就存在self.buffer这个数组中，每次reset把把归零

        
        # if len(self.record_bitratelevel) > 20:
        #     self.record_bitratelevel.pop(0)  # 大于20了就把一开始的剔除
        # # print(type(self.dataset.FovRecord[self.video_slot][0]))
        # # print(self.dataset.FovRecord[self.video_slot][0])

        # QoE
        # 码率  ## 用字典吧，直接加

        # 现在的逻辑是第一个FoV为1的当成主视角，也就是最高优先级 主视角0.8 主视角相邻的tile中 有1的给0.2，0的给0.1
        artificial_range_1s = {}
        range_temp = [0,1,2,3,4,5]
        main_perspective = 0
        main_near = []
        artificial_range_2s = {}
        for i in range_temp:
            if FoV[self.video_slot + 1][i].astype(int) == 1:
                # count = list(FoV[self.video_slot + 1].astype(int)).count(1)
                artificial_range_1s[i] = 1  # 主视野内权重为0.8
                range_temp.remove(i)
                main_perspective = i
                for j in FoV_Adjacent[i]:
                    range_temp.remove(j)
                    if FoV[self.video_slot + 1][j].astype(int) == 1:
                        artificial_range_1s[j] = 0.5
                        main_near.append(j)
                    else:
                        artificial_range_1s[j] = 0.2

            artificial_range_1s[range_temp[0]] = 0 ## ???

        bitrate_level_total = 0
        # 乘哪个FoV这个判断直接在state中考量，，等一下，有点问题
        for i in range(6):  # 6个面
            bitrate_level_total += action[i] * artificial_range_1s[i]

        # 帧内平滑度
        # 考虑FoV的东西，这个平滑度考虑主视角和其他四片的差异
        smoothness_within_frame = 0
        for i in FoV_Adjacent[main_perspective]:
            if i in main_near:
               smoothness_within_frame += abs(action[main_perspective] - action[i])
            else:
                smoothness_within_frame += abs(0.5 * (action[main_perspective] - action[i]))

        bitrate = 0
        for i in range(6):
            bitrate += action[i]

        alpha = 1.8
        beta = 0.3
        gamma = 0.5
        reward = alpha * bitrate_level_total - (6.5/18 * bitrate) - beta * smoothness_within_frame - rebuffer_event
        # reward = 10*reward
        # print(current_option,bitrate_level_total,smoothness_within_frame,frame_intra_smoothness,rebuffer_event,reward)
        
        # with open(DIR+'.txt',"a") as f:
        #     # print(current_option,bitrate_level_total,smoothness_within_frame,frame_intra_smoothness,rebuffer_event)
        #     f.write("reward: "+str(reward)+",  option值"+str(current_option)+",  比特率水平"+str(bitrate_level_total)+",  帧内平滑度"+str(smoothness_within_frame)+",  rebuffer事件"+str(rebuffer_record) + ",  最大不平滑度"+str(max_smooth)+"\n")
        with open(DIR+"_result.txt","a") as f:
            f.write("FoV2: "+str(FoV[self.video_slot+1])+",  因FoV做的决策: " + str(action)+ ",   option: " + str(current_option) +  "\n")

        self.video_slot += 1

        if self.video_slot >= len(self.dataset.throughput)-2:
            done = True
            self.video_slot = 0


        FoV_state = FoV[self.video_slot+1]
        for i in range(0,6):
            self.state[i] = FoV_state[i]
        latency = 0;data_amount = 0;##这俩从系统上测吧？
        _ = [bitrate_level_total,smoothness_within_frame,rebuffer_event, latency, data_amount]
        return self.state, reward, done, _

if __name__=="__main__":
    env = TiledVideo()
    env.seed(3)
    print("pause")
